from collections import OrderedDict
from batchgenerators.augmentations.utils import random_crop_2D_image_batched, pad_nd_image
import numpy as np
from batchgenerators.dataloading import SlimDataLoaderBase
from multiprocessing import Pool
from nnunet.paths import preprocessing_output_dir, cropped_output_dir
from batchgenerators.utilities.file_and_folder_operations import *
from nnunet.training.dataloading.dataset_loading import load_dataset
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import SimpleITK as sitk
from nnunet.preprocessing.preprocessing import resample_patient
from collections import Counter
from skimage.morphology import label
import cv2
from pathlib import Path
import shutil
from enum import Enum

def pad_case_data(case_all_data_donly, case_all_data_segonly,  crop_size = (128, 128, 128),
                  margin_size = (32, 32, 32), data_pad_value=200, seg_pad_value=0):
    d, h, w = case_all_data_donly.shape

    nd = int(np.ceil(float(d) / crop_size[0]))
    nh = int(np.ceil((float(h) / crop_size[1])))
    nw = int(np.ceil((float(w) / crop_size[2])))

    # pad_size = [[margin_size[0], nd * crop_size[0] - d + margin_size[0]],
    #             [margin_size[1], nh * crop_size[1] - h + margin_size[1]],
    #             [margin_size[2], nw * crop_size[2] - w + margin_size[2]]]
    pad_size = [[margin_size[0], margin_size[0]//2],
                [margin_size[1], margin_size[1]//2],
                [margin_size[2], margin_size[2]//2]]

    case_all_data_donly = np.pad(case_all_data_donly, pad_size, "constant", **{'constant_values': data_pad_value})

    case_all_data_segonly = np.pad(case_all_data_segonly, pad_size, 'constant', **{'constant_values': seg_pad_value})

    print(case_all_data_donly.shape, case_all_data_segonly.shape)

    return case_all_data_donly, case_all_data_segonly


def resample_and_normalize(intensityproperties, data, target_spacing, properties, seg=None, force_separate_z=None):
    """
    data and seg must already have been transposed by transpose_forward. properties are the un-transposed values
    (spacing etc)
    :param data:
    :param target_spacing:
    :param properties:
    :param seg:
    :param force_separate_z:
    :return:
    """

    # target_spacing is already transposed, properties["original_spacing"] is not so we need to transpose it!
    # data, seg are already transposed. Double check this using the properties
    original_spacing_transposed = np.array(properties["original_spacing"])[[0,1,2]]
    before = {
        'spacing': properties["original_spacing"],
        'spacing_transposed': original_spacing_transposed,
        'data.shape (data is transposed)': data.shape
    }
    data, seg = resample_patient(data, seg, np.array(original_spacing_transposed), target_spacing, 3, 1,
                                 force_separate_z=force_separate_z, order_z_data=0, order_z_seg=0)
    after = {
        'spacing': target_spacing,
        'data.shape (data is resampled)': data.shape
    }
    print("before:", before, "\nafter: ", after, "\n")

    if seg is not None:  # hippocampus 243 has one voxel with -2 as label. wtf?
        seg[seg < -1] = 0

    properties["size_after_resampling"] = data[0].shape
    properties["spacing_after_resampling"] = target_spacing
    use_nonzero_mask = False


    print("normalization...")

    for c in range(len(data)):
        # clip to lb and ub from train data foreground and use foreground mn and sd from training data
        assert intensityproperties is not None, "ERROR: if there is a CT then we need intensity properties"
        mean_intensity = intensityproperties[c]['mean']
        std_intensity = intensityproperties[c]['sd']
        lower_bound = intensityproperties[c]['percentile_00_5']
        upper_bound = intensityproperties[c]['percentile_99_5']
        #data[c] = np.clip(data[c], lower_bound, upper_bound)
        data[c] = np.clip(data[c], 200, upper_bound)

        #sitk.WriteImage(sitk.GetImageFromArray(data[c]), 'clip_data.nii')

        #data[c] = (data[c] - mean_intensity) / std_intensity

    print("normalization done")
    return data, seg, properties

def load_cropped(cropped_output_dir, case_identifier):
    all_data = np.load(os.path.join(cropped_output_dir, "%s.npz" % case_identifier))['data']
    data = all_data[:-1].astype(np.float32)
    seg = all_data[-1:]
    with open(os.path.join(cropped_output_dir, "%s.pkl" % case_identifier), 'rb') as f:
        properties = pickle.load(f)
    return data, seg, properties



def loadAndResampleData(cropped_output_dir, case_identifier, intensityproperties):
    transpose_forward = [0, 1, 2]
    target_spacing = np.array([1.0, 1.0, 1.0]).astype(np.float64)
    force_separate_z = None

    data, seg, properties = load_cropped(cropped_output_dir, case_identifier)
    data = data.transpose((0, *[i + 1 for i in transpose_forward]))
    seg = seg.transpose((0, *[i + 1 for i in transpose_forward]))

    # sitk.WriteImage(sitk.GetImageFromArray(data[0]), 'data.nii')
    # sitk.WriteImage(sitk.GetImageFromArray(seg[0]), 'seg.nii')

    data, seg, properties = resample_and_normalize(intensityproperties, data, target_spacing,
                                                        properties, seg, force_separate_z)

    return data, seg, properties




def save_patch_info_pkl(save_dir, case_identifier, patch_index, complement_label, pacth_label):
    info = OrderedDict()
    info['case_identifier'] = case_identifier
    info['ctvolume_name'] = "ctvolume_" + str(patch_index) + ".npy"
    info['complement_class'] = complement_label
    info['pacth_label'] = pacth_label
    info['segolume_name'] = "segolume_" + str(patch_index) + ".npy"
    write_pickle(info, os.path.join(save_dir, case_identifier + "_" + str(patch_index) + ".pkl"))


def analyze_seg_classes(mask):
    vol_per_voxel = np.prod([1.0,1.0,1.0])

    # ad 1)
    unique_classes = np.unique(mask)

    # 4) check if all in one region
    cur_all_classes = sorted(unique_classes.tolist())
    del cur_all_classes[0]

    volume_per_class = {}
    for c in cur_all_classes:
        volume_per_class[c] = np.sum(mask == c) * vol_per_voxel

    return volume_per_class


class PatchStatus(Enum):
    """
    Patch and mask intersection.
    """
    complemet = 0
    incomplement = 1
    containsnothing = 2


def generateIterativePatch(data, seg, case_identifier, save_patch_dir):
    crop_size = (128, 128, 128) #d h w
    LowerHU = 200.0

    x_crop_radius = crop_size[2] // 2
    y_crop_radius = crop_size[1] // 2
    z_crop_radius = crop_size[0] // 2

    x_step = int(x_crop_radius)
    y_step = int(y_crop_radius)
    z_step = int (z_crop_radius//2)

    segment_volume_pixels_thresh = 1000

    case_all_data_donly = data[0]
    case_all_data_segonly = seg[0]


    # case_all_data_donly = np.transpose(case_all_data_donly, [1, 2, 0])
    # case_all_data_segonly = np.transpose(case_all_data_segonly, [1, 2, 0])

    # 2. Pad data 左上角定点pad(128, 128, 128)
    case_all_data_segonly[case_all_data_segonly < 0] = 0
    # case_all_data_donly, case_all_data_segonly = pad_case_data(case_all_data_donly, case_all_data_segonly,
    #                                                            crop_size, crop_size,
    #                                                            data_pad_value=LowerHU, seg_pad_value=0)

    instace_memory_volume = np.zeros(case_all_data_segonly.shape, dtype=case_all_data_segonly.dtype)

    data_itk = sitk.GetImageFromArray(case_all_data_donly)
    # data_itk.SetSpacing(properties['itk_spacing'])
    # data_itk.SetOrigin(properties['itk_origin'])
    # data_itk.SetDirection(properties['itk_direction'])
    sitk.WriteImage(data_itk, 'clip_data_pad.nii')

    data_itk = sitk.GetImageFromArray(case_all_data_segonly)
    # data_itk.SetSpacing(properties['itk_spacing'])
    # data_itk.SetOrigin(properties['itk_origin'])
    # data_itk.SetDirection(properties['itk_direction'])
    sitk.WriteImage(data_itk, 'clip_mask_pad.nii')


    d, h, w = case_all_data_donly.shape
    seg_volume_per_class = analyze_seg_classes(case_all_data_segonly)


    def centerConvertzyxzyx(_centerzyx):
        _cropZYXZYX = np.array([_centerzyx[0] - z_crop_radius, _centerzyx[1] - y_crop_radius, _centerzyx[2] - x_crop_radius,
                               _centerzyx[0] + z_crop_radius - 1, _centerzyx[1] + y_crop_radius - 1,
                               _centerzyx[2] + x_crop_radius - 1])
        _cropZYXZYX[0::3] = np.clip(_cropZYXZYX[0::3], 0, d - 1)
        _cropZYXZYX[1::3] = np.clip(_cropZYXZYX[1::3], 0, h - 1)
        _cropZYXZYX[2::3] = np.clip(_cropZYXZYX[2::3], 0, w - 1)
        return _cropZYXZYX

    def crop_ct_seg(ct3D, seg3D, _centerzyx):
        _cropZYXZYX = centerConvertzyxzyx(_centerzyx)
        crop_ct = ct3D[_cropZYXZYX[0]: _cropZYXZYX[3] + 1, _cropZYXZYX[1]: _cropZYXZYX[4] + 1,
                  _cropZYXZYX[2]: _cropZYXZYX[5] + 1]
        crop_seg = seg3D[_cropZYXZYX[0]: _cropZYXZYX[3] + 1, _cropZYXZYX[1]: _cropZYXZYX[4] + 1,
                   _cropZYXZYX[2]: _cropZYXZYX[5] + 1]

        if crop_ct.shape != crop_size:
            pad_size = [[crop_size[0] - crop_ct.shape[0], 0],
                        [crop_size[1] - crop_ct.shape[1], 0],
                        [crop_size[2] - crop_ct.shape[2], 0]]
            crop_ct = np.pad(crop_ct, pad_size, "constant", **{'constant_values': LowerHU})

        if crop_seg.shape != crop_size:
            pad_size = [[crop_size[0] - crop_seg.shape[0], 0],
                        [crop_size[1] - crop_seg.shape[1], 0],
                        [crop_size[2] - crop_seg.shape[2], 0]]
            crop_seg = np.pad(crop_seg, pad_size, 'constant', **{'constant_values': 0})

        return crop_ct, crop_seg

    def compute3DVolume(volume3D, _centerzyx):
        """
        计算指定区域内像素的数量.
        :param data:3D volume
        :param di: z direction index
        :param centerPoint: (x, y)
        :return:volume
        """
        _cropZYXZYX = centerConvertzyxzyx(_centerzyx)

        segimage = volume3D[_cropZYXZYX[0]: _cropZYXZYX[3]+1, _cropZYXZYX[1]:_cropZYXZYX[4]+1, _cropZYXZYX[2]:_cropZYXZYX[5]+1]
        segimage[segimage < 0] = 0
        maskimgae = segimage>LowerHU
        voxels = maskimgae.sum()
        return voxels, _cropZYXZYX

    def computeDataSegComplement(seg3D, _centerzyx):
        """
        计算指定区域内像素的数量.
        :param data:3D volume
        :param di: z direction index
        :param centerPoint: (x, y)
        :return:volume
        """
        volume_thresh1 = 0.98
        volume_thresh2 = 0.75

        #1)convert zyx center into zyxzyx
        _cropZYXZYX = centerConvertzyxzyx(_centerzyx)


        #2) crop 当前区域的volume
        cur_crop_seg_image = np.zeros(seg3D.shape, dtype=seg3D.dtype)
        cur_crop_seg_image[_cropZYXZYX[0]: _cropZYXZYX[3] + 1, _cropZYXZYX[1]:_cropZYXZYX[4] + 1, _cropZYXZYX[2]:_cropZYXZYX[5] + 1] = \
            seg3D[_cropZYXZYX[0]: _cropZYXZYX[3] + 1, _cropZYXZYX[1]:_cropZYXZYX[4] + 1, _cropZYXZYX[2]:_cropZYXZYX[5] + 1]

        # sitk.WriteImage(sitk.GetImageFromArray(cur_crop_seg_image),
        #                 os.path.join(save_patch_dir, "cur_crop_seg_image.nii"))

        #3)计算crop seg volume里面所有的labels
        seglabels:dict = analyze_seg_classes(cur_crop_seg_image)

        # 1)判断存在1块脊椎还是多块， 找出最近脊椎块的位置
        # 2）移动到细的坐标点（只有一块的情况， 2块的情况）
        if len(seglabels.keys()) > 1:
            crop_label_lists = []
            for label in list(seglabels.keys()):
                label_volumes = seglabels[label]
                label_where = np.where(cur_crop_seg_image == label)
                z_mean = np.mean(label_where[0], axis=0)
                crop_label_lists.append({"l": label,"v": label_volumes, "mean_z":z_mean})

            crop_label_lists.sort(key=lambda k:(k.get('mean_z', 0)))
            first_label = crop_label_lists[0]["l"]
            mask_vloume_overlapping = seglabels[first_label] / seg_volume_per_class[first_label]

            if mask_vloume_overlapping  >=  volume_thresh1:
                #第一个label的区域完全满足要求, 则返回第二个label区域的mean coordinate
                small_where = np.where(cur_crop_seg_image == crop_label_lists[1]["l"])
                mean_where = np.mean(small_where, axis=1)
                return PatchStatus.complemet, np.where(seg3D == first_label), mean_where, first_label
            # elif volume_thresh1<  mask_vloume_overlapping  <  volume_thresh1:
            #     small_where = np.where(cur_crop_seg_image == crop_label_lists[0]["l"])
            #     mean_where = np.mean(small_where, axis=1)
            #     return False, np.where(seg3D == first_label), mean_where
            else:
                # 第一个label的区域不满足要求, 则返回第一个label区域的新的mean coordinate
                small_where = np.where(cur_crop_seg_image == crop_label_lists[0]["l"])
                mean_where = np.mean(small_where, axis=1)
                return PatchStatus.incomplement, np.where(seg3D == first_label), mean_where, None
        elif len(seglabels.keys()) == 1:
            label = list(seglabels.keys())[0]
            curVolume = seglabels[label]

            small_where = np.where(cur_crop_seg_image == label)
            mean_where = np.mean(small_where, axis=1)

            if curVolume >= seg_volume_per_class[label] * volume_thresh1:
                return PatchStatus.complemet, np.where(seg3D == label), mean_where, label
            else:
                return PatchStatus.incomplement, np.where(seg3D == label), mean_where, None
        return PatchStatus.containsnothing, None, None, None

    patch_index = 0
    start_z = 0 - z_crop_radius + z_crop_radius//2
    start_y = 0 - y_crop_radius + y_crop_radius//2
    start_x = 0 - x_crop_radius + x_crop_radius //2
    end_z = d
    change_z = True

    while (start_z < end_z):
        tmp_zyx = (start_z,  start_y, start_x)

        print('case_identifier: ', case_identifier, ' center zyx: ', tmp_zyx)
        if change_z:
            start_z += z_step

        curvoxels, _ = compute3DVolume(case_all_data_donly, tmp_zyx)
        if curvoxels >= segment_volume_pixels_thresh:
            complement_flag, label_coord, mean_where, patch_label = computeDataSegComplement(case_all_data_segonly, tmp_zyx)

            crop_ct_data, crop_seg_data = crop_ct_seg(case_all_data_donly, case_all_data_segonly, tmp_zyx)

            sitk.WriteImage(sitk.GetImageFromArray(crop_ct_data),
                            os.path.join(save_patch_dir, 'ctvolume_' + str(patch_index) + ".nii"))
            sitk.WriteImage(sitk.GetImageFromArray(crop_seg_data),
                            os.path.join(save_patch_dir, 'segvolume_' + str(patch_index) + ".nii"))

            if complement_flag == PatchStatus.complemet:
                #1)放入实例存储器
                instace_memory_volume[label_coord] = case_all_data_segonly[label_coord]
                #2)seg mask相关联的label区域清空
                case_all_data_segonly[label_coord] = 0

                # sitk.WriteImage(sitk.GetImageFromArray(instace_memory_volume),
                #                 os.path.join(save_patch_dir, "instace_memory_volume.nii"))
                #
                # sitk.WriteImage(sitk.GetImageFromArray(case_all_data_segonly),
                #                 os.path.join(save_patch_dir, "case_all_data_segonly.nii"))

                #3)ct data相关联的label区域清空
                case_all_data_donly[label_coord] = LowerHU
                print("index: ", patch_index, "Save Instance mem.")
                save_patch_info_pkl(save_patch_dir, case_identifier, patch_index, complement_flag, patch_label)
            elif complement_flag == PatchStatus.incomplement:
                print("index: ", patch_index, "Create incomplement patch.")
                save_patch_info_pkl(save_patch_dir, case_identifier, patch_index, complement_flag, patch_label)
            elif complement_flag == PatchStatus.containsnothing:
                print("index: ", patch_index, "Create containsnothing patch.")
                save_patch_info_pkl(save_patch_dir, case_identifier, patch_index, complement_flag, patch_label)

                # x y direction
                change_z = False
                start_y = (start_y + y_step)
                start_x = (start_x + x_step)
                if start_y >= h - y_crop_radius or start_x >=  w - x_crop_radius:
                    start_y = start_y % (h - y_crop_radius)
                    start_x = start_x % (w - x_crop_radius)
                    change_z = True

            if mean_where is not None:
                start_z, start_y, start_x = (mean_where.astype(np.int)).tolist()
            patch_index += 1

        else:
            change_z = False
            start_y = (start_y + y_step)
            start_x = (start_x + x_step)
            if start_y >= h - y_crop_radius or start_x >= w - x_crop_radius:
                change_z = True
                start_y = start_y % (h - y_crop_radius)
                start_x = start_x % (w - x_crop_radius)


    sitk.WriteImage(sitk.GetImageFromArray(instace_memory_volume),
                    os.path.join(save_patch_dir, "instace_memory_volume.nii"))
    sitk.WriteImage(sitk.GetImageFromArray(case_all_data_segonly),
                    os.path.join(save_patch_dir, "case_all_data_segonly.nii"))

    #print("compare instace_memory and seg mask data: ", np.sum(np.logical_xor(instace_memory_volume, case_all_data_segonly)))
    print("compare data_segonly: ",np.sum(case_all_data_segonly))




def generate_case_identifier(cropped_output_dir, case_identifier, intensityproperties, save_patch_dir):
    print(case_identifier)
    data, seg, properties = loadAndResampleData(cropped_output_dir, case_identifier, intensityproperties)
    save_patch_dir = os.path.join(save_patch_dir, case_identifier)

    #if not os.path.exists(save_patch_dir):
    maybe_mkdir_p(save_patch_dir)
    generateIterativePatch(data, seg, case_identifier, save_patch_dir)



def test_generateFunc():
    t = "VerSe19"
    p = join(cropped_output_dir, t)
    # plans_fname = join(preprocessing_output_dir, t, "nnUNetPlans_plans_3D.pkl")
    # dataset = load_dataset(p)

    with open(join(p, "intensityproperties.pkl"), 'rb') as f:
        intensityproperties = pickle.load(f)

    #train_verse006 train_verse051
    case_identifier = 'train_verse051'

    save_patch_dir = join(preprocessing_output_dir, t, "patch")
    maybe_mkdir_p(save_patch_dir)

    generate_case_identifier(p, case_identifier, intensityproperties, save_patch_dir)



def generate_patches():
    t = "VerSe19"
    p = join(cropped_output_dir, t)
    with open(join(p, "intensityproperties.pkl"), 'rb') as f:
        intensityproperties = pickle.load(f)

    case_identifier_list = sorted(Path(p).glob("*.npz"))
    case_identifier_list = [identifier.name.split(".")[0] for identifier in case_identifier_list]

    save_patch_dir = join(preprocessing_output_dir, t, "patch")
    maybe_mkdir_p(save_patch_dir)

    for case_identifier in case_identifier_list:
        generate_case_identifier(p, case_identifier, intensityproperties, save_patch_dir)




def test_itk():
    p = '/home/blake/data/medical/datasets/vertebral/VerSe_2019/training_phase_1_release'
    case_identifier_list = sorted(Path(p).glob("*_seg_n.nii"))
    case_identifier_list = [identifier.name.split("_")[0] for identifier in case_identifier_list]
    print(len(case_identifier_list))
    # fix_d = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)
    for case_identifier in case_identifier_list:
        ct_f = os.path.join(p, case_identifier+'_n.nii')
        seg_f = os.path.join(p, case_identifier+'_seg_n.nii')
        os.rename(ct_f, os.path.join(p, case_identifier+'.nii'))
        os.rename(seg_f, os.path.join(p, case_identifier + '_seg.nii'))

        # fitk = sitk.ReadImage(ct_f)
        # print(case_identifier, ": before ", fitk.GetSpacing(), sitk.GetArrayFromImage(fitk).shape)
        # img = sitk.GetArrayFromImage(fitk)
        # img = np.transpose(img, [1, 2, 0])
        # spcing = fitk.GetSpacing()
        #
        # oitk = sitk.GetImageFromArray(img)
        # oitk.SetSpacing(spcing)
        # print(case_identifier, ": after ", oitk.GetSpacing(), sitk.GetArrayFromImage(oitk).shape)
        # ver_ctf = str(Path(ct_f).with_name((Path(ct_f).stem+ "_n.nii")))
        # sitk.WriteImage(oitk, ver_ctf)
        #
        # fitk = sitk.ReadImage(seg_f)
        # print(case_identifier, ": before ", fitk.GetSpacing(), sitk.GetArrayFromImage(fitk).shape)
        # img = sitk.GetArrayFromImage(fitk)
        # img = np.transpose(img, [1, 2, 0])
        # oitk = sitk.GetImageFromArray(img)
        #
        # print(case_identifier, ": after ", oitk.GetSpacing(), sitk.GetArrayFromImage(oitk).shape)
        # ver_ctf = str(Path(seg_f).with_name((Path(seg_f).stem + "_n.nii")))
        # sitk.WriteImage(oitk, ver_ctf)



def main():
    print("VerseGenerate...")
    test_generateFunc()

if __name__ == '__main__':
    main()
